"""
database for image search

file_md5, milvus_id, image_path
# include file md5 to prevent duplicate images
"""
import os
import pymysql
import random
import sys
from engine.utils.logs import LOGGER
from engine.utils import get_config
from engine import file_md5


class DBImageSearch(object):
    def __init__(self, config):
        self.config = config.MYSQL
        self.create_connection()
    
    def create_connection(self):
        self.conn = pymysql.connect(
            host=self.config.HOST, 
            port=self.config.PORT,
            user=self.config.USER,
            password=self.config.PWD,
            database=self.config.DB,
            local_infile=True)
        self.cursor = self.conn.cursor()
    
    def test_connection(self):
        try:
            self.conn.ping()
        except Exception:
            self.create_connection()
    
    def create_table(self, table_name):
        self.test_connection()
        sql = "create table if not exists " + table_name + "(file_md5 TEXT, milvus_id TEXT, image_path TEXT);"
        try:
            self.cursor.execute(sql)
            LOGGER.debug(f"MYSQL create table: {table_name} with sql: {sql}")
        except Exception as e:
            LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
            sys.exit(1)
    
    def exists(self, table_name, file_md5):
        self.test_connection()
        sql = "select count(*) from " + table_name + " where file_md5 = '" + file_md5 + "';"
        try:
            self.cursor.execute(sql)
            results = self.cursor.fetchall()
            LOGGER.debug(f"MYSQL exists table:{table_name}")
            return results[0][0]
        except Exception as e:
            LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
            sys.exit(1)
    
    def insert(self, table_name, data):
        self.test_connection()
        sql = "insert into " + table_name + " (file_md5,milvus_id,image_path) values (%s,%s,%s);"
        try:
            # self.cursor.executemany(sql, data)
            self.cursor.execute(sql, data)
            self.conn.commit()
            LOGGER.debug(f"MYSQL loads data to table: {table_name} successfully")
        except Exception as e:
            LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
            sys.exit(1)
    
    def insert_if_not_exists(self, table_name, data):
        file_md5, milvus_id, image_path = data
        if not self.exists(table_name, file_md5):
            self.insert(table_name, data)
        else:
            # do logging
            pass
    
    def search_by_milvus_id(self, ids, table_name):
        self.test_connection()
        str_ids = str(ids).replace('[', '').replace(']', '')
        sql = "select image_path from " + table_name + " where milvus_id in (" + str_ids + ") order by field (milvus_id," + str_ids + ");"
        try:
            self.cursor.execute(sql)
            results = self.cursor.fetchall()
            results = [res[0] for res in results]
            LOGGER.debug("MYSQL search by milvus id.")
            return results
        except Exception as e:
            LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
            sys.exit(1)

    def search_by_file_md5(self, md5s, table_name):
        self.test_connection()
        str_ids = str(md5s).replace('[', '').replace(']', '')
        sql = "select image_path from " + table_name + " where file_md5 in (" + str_ids + ") order by field (file_md5," + str_ids + ");"
        try:
            self.cursor.execute(sql)
            results = self.cursor.fetchall()
            results = [res[0] for res in results]
            LOGGER.debug("MYSQL search by file md5.")
            return results
        except Exception as e:
            LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
            sys.exit(1)

    def count(self, table_name):
        self.test_connection()
        sql = "select count(milvus_id) from " + table_name + ";"
        try:
            self.cursor.execute(sql)
            results = self.cursor.fetchall()
            LOGGER.debug(f"MYSQL count table:{table_name}")
            return results[0][0]
        except Exception as e:
            LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
            sys.exit(1)
    
    def delete_table(self, table_name):
        # Delete mysql table if exists
        self.test_connection()
        sql = "drop table if exists " + table_name + ";"
        try:
            self.cursor.execute(sql)
            LOGGER.debug(f"MYSQL delete table:{table_name}")
        except Exception as e:
            LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
            sys.exit(1)

    def delete_all_data(self, table_name):
        # Delete all the data in mysql table
        self.test_connection()
        sql = 'delete from ' + table_name + ';'
        try:
            self.cursor.execute(sql)
            self.conn.commit()
            LOGGER.debug(f"MYSQL delete all data in table:{table_name}")
        except Exception as e:
            LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
            sys.exit(1)


def test_mysql():
    config_path = 'configs.yaml'
    cfg = get_config(config_path)
    db = DBImageSearch(cfg)
    table_name = 'image_search'

    # create table
    db.create_table(table_name)

    c = db.count(table_name)
    print(c)

    image_folder = '/rinoshinme/data/image_search'
    for name in os.listdir(image_folder):
        print('inserting ', name)
        file_path = os.path.join(image_folder, name)
        md5_value = file_md5(file_path)
        milvus_id = str(random.randint(0, 999999)).zfill(6)
        item = (md5_value, milvus_id, file_path)
        db.insert_if_not_exists(table_name, item)
        # ret = db.exists(table_name, md5_value)
        # ret = db.search_by_file_md5([md5_value], table_name)
        # print(ret)

        break
